package client.block;

import client.texture.UV;
import common.block.AbstractBlock;
import common.util.ILogger;
import org.lwjgl.opengl.GL;
import org.lwjgl.stb.STBImage;

import javax.annotation.Nullable;
import java.nio.ByteBuffer;
import java.util.*;

import static org.lwjgl.opengl.GL46C.*;

public class MultiBlockAtlasBuilder {

    private final String MISSING_PATH = "";

    private final int TOTAL_W = 2048, TOTAL_H = 2048, PER_SIZE = 64;
    private final int LINE_COUNT = TOTAL_W / PER_SIZE, TOTAL_COUNT = LINE_COUNT * LINE_COUNT;

    private final int levelCount;
    // for every atlas, the first image is missing-block
    private int count = 1;

    private final List<TextureCollector> collectors = new ArrayList<>();
    private TextureCollector currentCollector = null;

    private int[] textureIds = new int[0];
    private final int buffer;

    public MultiBlockAtlasBuilder(int levelCount) {
        this.levelCount = levelCount;
        buffer = glGenBuffers();
        glBindBuffer(GL_PIXEL_UNPACK_BUFFER, buffer);
        if (GL.getCapabilities().GL_ARB_buffer_storage) {
            glBufferStorage(GL_PIXEL_UNPACK_BUFFER, PER_SIZE * PER_SIZE * 4, GL_DYNAMIC_STORAGE_BIT);
        } else {
            glBufferData(GL_PIXEL_UNPACK_BUFFER, PER_SIZE * PER_SIZE * 4, GL_DYNAMIC_COPY);
        }
    }

    public void put(AbstractBlock block, String name, String texture) {
        if (currentCollector == null) {
            currentCollector = new TextureCollector(block);
            collectors.add(currentCollector);
        }
        if (currentCollector.block == block) {
            currentCollector.textures.put(name, texture);
            count++;
        } else {
            TextureCollector next = null;
            if (count > TOTAL_COUNT) {
                count -= currentCollector.textures.size();
                next = currentCollector;
                collectors.remove(currentCollector);
            }
            buildAtlas();
            collectors.clear();
            count = 0;
            if (next != null) {
                collectors.add(next);
                count += next.textures.size();
            }
            currentCollector = new TextureCollector(block);
            collectors.add(currentCollector);
        }
    }

    public int[] end() {
        if (!collectors.isEmpty()) {
            buildAtlas();
        }
        collectors.clear();
        glBindBuffer(GL_PIXEL_UNPACK_BUFFER, 0);
        glDeleteBuffers(buffer);
        return textureIds;
    }

    private void buildAtlas() {
        int texture = glGenTextures();
        glBindTexture(GL_TEXTURE_2D, texture);
        int[] newIds = Arrays.copyOf(textureIds, textureIds.length + 1);
        newIds[textureIds.length] = texture;
        textureIds = newIds;

        if (GL.getCapabilities().GL_ARB_texture_storage) {
            glTexStorage2D(GL_TEXTURE_2D, levelCount, GL_RGBA4, TOTAL_W, TOTAL_H);
        } else {
            glTexImage2D(GL_TEXTURE_2D, 0, GL_RGBA4, TOTAL_W, TOTAL_H, 0, GL_RGBA, GL_UNSIGNED_BYTE, 0);
        }
        int x = 1, y = 0;
        for (TextureCollector collector : collectors) {
            BlockTextureAtlas atlas = new BlockTextureAtlas(texture);
            loadImage(MISSING_PATH, "missing", 0, 0);
            for (Map.Entry<String, String> entry : collector.textures.entrySet()) {
                String textureName = entry.getKey();
                String texturePath = entry.getValue();
                UV uv = loadImage(texturePath, textureName, x, y);
                if (uv != null) {
                    x++;
                    if (x >= LINE_COUNT) {
                        x = 0;
                        y++;
                    }
                }
                atlas.put(textureName, uv);
            }
            AbstractBlock block = collector.block;
            block.setTextureAtlas(atlas);
            for (int id : textureIds) {
                if (id == atlas.textureId) return;
            }
        }
        glBindTexture(GL_TEXTURE_2D, 0);
    }

    @Nullable
    private UV loadImage(String texturePath, String textureName, int x, int y) {
        int[] width = new int[1], height = new int[1], channels = new int[1];
        ByteBuffer buffer = STBImage.stbi_load(texturePath, width, height, channels, 4);
        if (buffer == null) {
            ILogger.getInstance().error("Can't load image " + textureName + ": " + texturePath);
            return null;
        }
        if (width[0] > PER_SIZE || height[0] >= PER_SIZE) {
            ILogger.getInstance().error("Image is too large: " + width[0] + " X " + height[0] + ", max is " + PER_SIZE + " X " + PER_SIZE);
            return null;
        }
        UV uv = new UV(width[0], height[0], x * PER_SIZE, y * PER_SIZE, TOTAL_W, TOTAL_H);
        glBufferSubData(GL_PIXEL_UNPACK_BUFFER, 0, buffer);
        glTexSubImage2D(GL_TEXTURE_2D, 0, uv.x0, uv.y0, uv.width, uv.height, GL_RGBA, GL_UNSIGNED_BYTE, 0L);
        return uv;
    }

    static class TextureCollector {
        protected AbstractBlock block;
        protected Map<String, String> textures = new HashMap<>();

        public TextureCollector(AbstractBlock block) {
            this.block = block;
        }
    }
}
